import pickle
import matplotlib.pyplot as plt

# COCO17骨架连接
skeleton = [
    [15, 13], [13, 11], [16, 14], [14, 12],
    [11, 12], [5, 11], [6, 12], [5, 6],
    [5, 7], [6, 8], [7, 9], [8, 10],
    [1, 2], [0, 1], [0, 2], [1, 3], [2, 4],
    [3, 5], [4, 6]
]

# 每个骨架线的颜色
colors = [
    'red', 'red', 'green', 'green',
    'purple', 'orange', 'orange', 'blue',
    'orange', 'orange', 'pink', 'pink',
    'gray', 'gray', 'gray', 'gray', 'gray',
    'cyan', 'cyan'
]

def load_keypoints_from_pkl(path):
    with open(path, 'rb') as f:
        data = pickle.load(f)
    return data  # shape: [frames, 17, 3]

def draw_pose(ax, keypoints, skeleton, colors, frame_idx):
    for (pt1, pt2), color in zip(skeleton, colors):
        x1, y1 = keypoints[pt1][:2]
        x2, y2 = keypoints[pt2][:2]
        ax.plot([x1, x2], [y1, y2], color=color, linewidth=3)

    for i, kp in enumerate(keypoints):
        x, y = kp[:2]
        ax.plot(x, y, 'o', color='black', markersize=4)
        ax.text(x + 2, y, str(i), fontsize=7, color='blue')  # 显示关键点编号

    # ax.set_title(f"Frame {frame_idx}")
    ax.set_xlim(0, 250)
    ax.set_ylim(400, 0)
    ax.set_aspect('equal')
    ax.axis('off')

def show_frames(sequence):
    for i, keypoints in enumerate(sequence):
        fig, ax = plt.subplots(figsize=(6, 10))
        draw_pose(ax, keypoints, skeleton, colors, i)
        plt.show()  # 手动关闭一帧窗口后再显示下一帧


# def show_frames(sequence):
#     for i, keypoints in enumerate(sequence):
#         fig, ax = plt.subplots(figsize=(6, 10))  # 原始画布大小
#         draw_pose(ax, keypoints, skeleton, colors, i)
#         plt.savefig(f"frame_{i:03d}.png", bbox_inches='tight', pad_inches=0.05)
#         plt.show()  # 不弹窗，只保存


# 替换为你的 pkl 文件路径
pkl_path = r"E:\datasets\2.Gait3D\2D_Poses-pkl\0030\camid7_videoid1\seq0\seq0.pkl"
sequence = load_keypoints_from_pkl(pkl_path)
show_frames(sequence)
